/*
 * Copyright 2022 The Open Islands Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.openislands.oi.scheduling;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.openislands.oi.util.CommonUtils;
import org.openislands.oi.util.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.yaml.snakeyaml.Yaml;
import org.yaml.snakeyaml.constructor.Constructor;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.InputStream;
import java.util.*;

public class DAG {
    private static Logger logger = LoggerFactory.getLogger(DAG.class);

    private DAGRole role;
    Map<String, Object> resource;
    private List<DAGOperator> operator;

    private Map<String, Integer> role2Index = new HashMap<>();              //  计算角色name->index的映射关系
    private Map<String, Integer> name2Index = new HashMap<>();              // 计算算子topology：保存算子name->index的映射关系
    private Map<String, List<String>> upstream = new HashMap<>();           // 保存算子上游输入
    private Map<String, List<String>> downstream = new HashMap<>();         // 保存算子下游输出

    public static DAG newDAG(String config){
        Yaml yaml = new Yaml(new Constructor(DAG.class));
        return yaml.load(config);
    }

    public DAGRole getRole() {
        return role;
    }

    public void setRole(DAGRole role) {
        this.role = role;
    }

    public Map<String, Object> getResource() {
        return resource;
    }

    public void setResource(Map<String, Object> resource) {
        this.resource = resource;
    }

    public List<DAGOperator> getOperator() {
        return operator;
    }

    public void setOperator(List<DAGOperator> operator) {
        this.operator = operator;
    }

    public Map<String, Integer> getRole2Index() {return role2Index;}
    public Map<String, Integer> getName2Index() {return name2Index;}
    public Map<String, List<String>> getUpstream() {return upstream;}
    public Map<String, List<String>> getDownstream() {return downstream;}

    // 检查dag配置的合法性; 更新算子resource属性；保存算子name->index映射(NodeID在不同角色中可以重复)
    public void checkRole() throws DAGException, JsonProcessingException {
        // 检查role是否为空
        if (role == null) {
            DAGErrorEnums error = DAGErrorEnums.ROLE_MAPPING_EMPTY;
            logger.error("error={}", error.getMessage());
            throw new DAGException(error);
        }
        List<String> guest = role.getGuest(), host = role.getHost(), arbiter = role.getArbiter();
        if ((guest == null || guest.isEmpty() == true) &&
                (host == null || host.isEmpty() == true) &&
                (arbiter == null || arbiter.isEmpty() == true)) {
            DAGErrorEnums error = DAGErrorEnums.ROLE_MAPPING_EMPTY;
            logger.error("error={}", error.getMessage());
            throw new DAGException(error);
        } else {
            int index = 0;
            if (guest != null) {
                for (String key : guest) {
                    if (role2Index.containsKey(key)) {
                        DAGErrorEnums error = DAGErrorEnums.ROLE_ID_DUPLICATE;
                        logger.info("id={}, error={}", key, error.getMessage());
                    } else {
                        role2Index.put(key, index);
                        index++;
                    }
                }
            }
            if (host != null) {
                for (String key : host) {
                    if (role2Index.containsKey(key)) {
                        DAGErrorEnums error = DAGErrorEnums.ROLE_ID_DUPLICATE;
                        logger.info("id={}, error={}", key, error.getMessage());
                    } else {
                        role2Index.put(key, index);
                        index++;
                    }
                }
            }
            if (arbiter != null) {
                for (String key : arbiter) {
                    if (role2Index.containsKey(key)) {
                        DAGErrorEnums error = DAGErrorEnums.ROLE_ID_DUPLICATE;
                        logger.info("id={}, error={}", key, error.getMessage());
                    } else {
                        role2Index.put(key, index);
                        index++;
                    }
                }
            }
        }
        logger.info("role2Index={}", role2Index);
        return;
    }
    public void checkOperator() throws DAGException, JsonProcessingException {
        // 检查算子列表是否为空
        if (operator == null || operator.size() == 0) {
            DAGErrorEnums error = DAGErrorEnums.OPERATOR_LIST_EMPTY;
            logger.error("error={}", error.getMessage());
            throw new DAGException(error);
        }

        // 遍历算子列表
        for (int index = 0; index < operator.size(); index++) {
            // 保存算子name是否唯一；并保存name->index的映射关系
            DAGOperator op = operator.get(index);
            String name = op.getName();
            if (name2Index.containsKey(name) == true) {
                DAGErrorEnums error = DAGErrorEnums.OPERATOR_NAME_DUPLICATE;
                logger.error("name={}, error={}", name, error.getMessage());
                throw new DAGException(error);
            } else {
                name2Index.put(name, index);
            }

            // 检查算子name是否合法: name必须是数字、字母、下划线
            if (name == null || name.isEmpty()) {
                DAGErrorEnums error = DAGErrorEnums.OPERATOR_NAME_NOT_EXISTS;
                logger.error("name={}, error={}", name, error.getMessage());
                throw new DAGException(error);
            }
            if (StringUtils.isLetterOrDigit(name, "_-") == false) {
                DAGErrorEnums error = DAGErrorEnums.OPERATOR_NAME_INVALID;
                logger.error("name={}, error={}", name, error.getMessage());
                throw new DAGException(error);
            }

            // 检查算子module是否合法: module必须配置
            String module = op.getModule();
            if (module == null || module.isEmpty()) {
                DAGErrorEnums error = DAGErrorEnums.OPERATOR_MODULE_NOT_EXISTS;
                logger.error("{}.module={}, error={}", name, module, error.getMessage());
                throw new DAGException(error);
            }

            // 算子param字段配置可选，跳过

            // 合并更新算子resource属性
            Map<String, Object> merge = CommonUtils.merge(resource, op.getResource());
            op.setResource(merge);

            // 检查算子input是否存在: 大部分算子input应该配置；只有最前面承接输入的算子没有input
            DAGOperatorInput input = op.getInput();
            if (input == null) {
                DAGErrorEnums error = DAGErrorEnums.OPERATOR_INPUT_NOT_EXISTS;
                logger.warn("{}.input=null, error={}", name, error.getMessage());
                // throw new DAGException(error);
            }

            // 检查算子output是否存在: output没有配置则给出提示
            DAGOperatorOutput output = op.getOutput();
            if (output == null) {
                DAGErrorEnums error = DAGErrorEnums.OPERATOR_OUTPUT_NOT_EXISTS;
                logger.warn("{}.output=null, error={}", name, error.getMessage());
                // throw new DAGException(error);
            } else {
                List<String> outputData = output.getData();
                List<String> outputModel = output.getModel();
                // 检查output.data/model是否同时为空
                if (outputData == null && outputModel == null) {
                    DAGErrorEnums error = DAGErrorEnums.OPERATOR_OUTPUT_ALL_NULL;
                    logger.error("{}.output.data=null&&{}.output.model=null, error={}", name, name, error.getMessage());
                    throw new DAGException(error);
                }
                // 检查output.data类型是否为String
                if (outputData != null && outputData.isEmpty() == false) {
                    if (outputData.get(0) instanceof String == false) {
                        DAGErrorEnums error = DAGErrorEnums.OPERATOR_OUTPUT_DATA_TYPE_ERROR;
                        logger.error("{}.output.data!=string, error={}", name, error.getMessage());
                        throw new DAGException(error);
                    }
                }
                // 检查output.model类型是否为String
                if (outputModel != null && outputModel.isEmpty() == false) {
                    if (outputModel.get(0) instanceof String == false) {
                        DAGErrorEnums error = DAGErrorEnums.OPERATOR_OUTPUT_MODEL_TYPE_ERROR;
                        logger.error("{}.output.model!=string, error={}", name, error.getMessage());
                        throw new DAGException(error);
                    }
                }
            }
        }
        return;
    }

    public void check() throws DAGException, JsonProcessingException {
        checkRole();;
        checkOperator();
        return;
    }

    // 计算dag图中的上下游依赖: operator.input可以看出依赖关系
    public void calDependency() throws DAGException {
        // 遍历dag图中的算子
        for (DAGOperator op : operator) {
            String name = op.getName();
            DAGOperatorInput input = op.getInput();
            if (input == null) {
                continue;
            }
            // 检查 operator.input.data部分依赖
            Map<String, List<String>> data = input.getData();
            if (data != null) {
                for (Map.Entry<String, List<String>> entry : data.entrySet()) {
                    for (String key : entry.getValue()) {
                        String[] elem = key.split("\\.");
                        if (elem.length != 2) {
                            DAGErrorEnums error = DAGErrorEnums.OPERATOR_INPUT_DATA_FORMAT_ERROR;
                            logger.error("{} <-- {}, error={}", name, key, error.getMessage());
                            throw new DAGException(error);
                        }
                        String prevName = elem[0], prevData = elem[1];
                        // 检查上游算子：名字是否合法；output.data是否包含当前data名
                        if (name2Index.containsKey(prevName) == false) {
                            DAGErrorEnums error = DAGErrorEnums.OPERATOR_INPUT_PREV_NAME_DISCONNECT;
                            logger.error("{} <-- {}, error={}", name, prevName, error.getMessage());
                            throw new DAGException(error);
                        }
                        // 检查上游算子output是否存在；output.data是否包含当前prevData
                        DAGOperator prevOp = operator.get(name2Index.get(prevName));
                        DAGOperatorOutput prevOutput = prevOp.getOutput();
                        boolean flag = true;
                        if (prevOutput == null) {
                            flag = false;
                        } else {
                            List<String> prevOutputData = prevOutput.getData();
                            if (prevOutputData == null || prevOutputData.isEmpty() == true) {
                                flag = false;
                            } else {
                                if (CommonUtils.contains(prevOutputData, prevData) == false) {
                                    flag = false;
                                } else {
                                    flag = true;
                                }
                            }
                        }

                        // 算子data部分连接不合法则报错；否则更新上下游依赖
                        if (flag == false) {
                            DAGErrorEnums error = DAGErrorEnums.OPERATOR_INPUT_PREV_DATA_DISCONNECT;
                            logger.error("{} <-- {}, error={}", name, prevName, error.getMessage());
                            throw new DAGException(error);
                        } else {
                            // 更新当前算子 的 上游依赖
                            if (upstream.containsKey(name) == false) {
                                List<String> dummy = upstream.getOrDefault(name, new ArrayList<>());
                                upstream.put(name, dummy);
                            }
                            List<String> upstreamMapping = upstream.get(name);
                            if (upstreamMapping.contains(prevName) == false) {
                                upstreamMapping.add(prevName);
                            }
                            // 更新前序算子 的 下游依赖
                            if (downstream.containsKey(prevName) == false) {
                                List<String> dummy = downstream.getOrDefault(prevName, new ArrayList<>());
                                downstream.put(prevName, dummy);
                            }
                            List<String> downstreamMapping = downstream.get(prevName);
                            if (downstreamMapping.contains(name) == false) {
                                downstreamMapping.add(name);
                            }
                        }
                    }
                }
            }

            // 检查 operator.input.model部分依赖
            Map<String, List<String>> model = input.getModel();
            if (model != null) {
                for (Map.Entry<String, List<String>> entry : model.entrySet()) {
                    for (String key : entry.getValue()) {
                        String[] elem = key.split("\\.");
                        // logger.info("elem={}", Arrays.toString(elem));
                        if (elem.length != 2) {
                            DAGErrorEnums error = DAGErrorEnums.OPERATOR_INPUT_MODEL_FORMAT_ERROR;
                            logger.error("{} <-- {}, error={}", name, key, error.getMessage());
                            throw new DAGException(error);
                        }
                        String prevName = elem[0], prevModel = elem[1];
                        // 检查上游算子：名字是否合法；output.data是否包含当前data名
                        if (name2Index.containsKey(prevName) == false) {
                            DAGErrorEnums error = DAGErrorEnums.OPERATOR_INPUT_PREV_NAME_DISCONNECT;
                            logger.error("{} <-- {}, error={}", name, prevName, error.getMessage());
                            throw new DAGException(error);
                        }
                        // 检查上游算子output是否存在；output.data是否包含当前prevData
                        DAGOperator prevOp = operator.get(name2Index.get(prevName));
                        DAGOperatorOutput prevOutput = prevOp.getOutput();
                        boolean flag = true;
                        if (prevOutput == null) {
                            flag = false;
                        } else {
                            List<String> prevOutputModel = prevOutput.getModel();
                            if (prevOutputModel == null || prevOutputModel.isEmpty() == true) {
                                flag = false;
                            } else {
                                if (CommonUtils.contains(prevOutputModel, prevModel) == false) {
                                    flag = false;
                                } else {
                                    flag = true;
                                }
                            }
                        }

                        // 算子model部分连接不合法则报错；否则更新上下游依赖
                        if (flag == false) {
                            DAGErrorEnums error = DAGErrorEnums.OPERATOR_INPUT_PREV_MODEL_DISCONNECT;
                            logger.error("{} <-- {}, error={}", name, prevName, error.getMessage());
                            throw new DAGException(error);
                        } else {
                            // 更新当前算子 的 上游依赖
                            if (upstream.containsKey(name) == false) {
                                List<String> dummy = upstream.getOrDefault(name, new ArrayList<>());
                                upstream.put(name, dummy);
                            }
                            List<String> upstreamMapping = upstream.get(name);
                            if (upstreamMapping.contains(prevName) == false) {
                                upstreamMapping.add(prevName);
                            }
                            // 更新前序算子 的 下游依赖
                            if (downstream.containsKey(prevName) == false) {
                                List<String> dummy = downstream.getOrDefault(prevName, new ArrayList<>());
                                downstream.put(prevName, dummy);
                            }
                            List<String> downstreamMapping = downstream.get(prevName);
                            if (downstreamMapping.contains(name) == false) {
                                downstreamMapping.add(name);
                            }
                        }
                    }
                }
            }
        }
        return;
    }

    // 获取DAG图中的root节点
    public List<DAGOperator> filterOperatorRoot() {
        List<DAGOperator> ops = new ArrayList<DAGOperator>();
        for (DAGOperator op : operator) {
            String name = op.getName();
            String module = op.getModule();
            DAGOperatorInput opInput = op.getInput();
            // 算子无input的挑选出来
            if (opInput == null) {
                ops.add(op);
                continue;
            }
        }
        return ops;
    }

    // 获取DAG图中指定类型的算子列表；比如获取图中Reader算子表示的所有输入
    public List<DAGOperator> filterOperatorByModule(String m) {
        if (m == null || m.isEmpty() == true) {
            logger.warn("filterOperatorByModule module=null");
            return null;
        }
        List<DAGOperator> ops = new ArrayList<DAGOperator>();
        for (DAGOperator op : operator) {
            String name = op.getName();
            String module = op.getModule();
            DAGOperatorInput opInput = op.getInput();
            // 算子.module指定类型也选出
            if (module.equals(m)) {
                ops.add(op);
            }
        }
        return ops;
    }

    // 获取指定名字的算子; 名字是算子唯一标识理论上应该只有一个满足
    public DAGOperator filterOperatorByName(String n) {
        // 输入算子名不能为空
        if (n == null || n.isEmpty() == true) {
            logger.warn("filterOperatorByName name=null");
            return null;
        }
        for (DAGOperator op : operator) {
            String name = op.getName();
            if (name.equals(n)) {
                return op;
            }
        }
        return null;
    }

    // 输入三元组算子名+角色+nodeID；过滤指定role元素精简算子属性
    public DAGOperator filterOperatorByTriple(String n, String role, String nodeId) throws JsonProcessingException {
        Set<String> roles = new HashSet<>(Arrays.asList("guest", "host", "arbiter"));
        if ((n == null || n.isEmpty() == true)) {
            logger.warn("filterOperator name=null");
            return null;
        }
        if ((role == null || role.isEmpty() == true)) {
            logger.warn("filterOperator role=null");
            return null;
        } else if (roles.contains(role) == false) {
            logger.warn("filterOperator role={} not in {}", role, roles);
            return null;
        }
        if ((nodeId == null || nodeId.isEmpty() == true)) {
            logger.warn("filterOperator nodeId=null");
            return null;
        } else if (role2Index.containsKey(nodeId) == false) {
            logger.warn("filterOperator nodeId={} not in dag graph", nodeId);
            return null;
        }

        // 待删除的角色集合
        roles.remove(role);
        logger.info("roles={}", roles);

        for (DAGOperator op : operator) {
            // 过滤算子名
            if (n.equals(op.getName()) == false) {
                continue;
            }

            // 过滤resource中的角色
            DAGOperator newOp = DAGOperator.fromJson(op.toJson());
            if (newOp.getResource() != null) {
                Map<String, Object> opResource = newOp.getResource();
                for (String r : roles) {
                    if (opResource.containsKey(r)) {
                        opResource.remove(r);
                    }
                }
                newOp.setResource(opResource);
            }

            // 过滤param中的角色
            if (newOp.getParam() != null) {
                Map<String, Object> opParam = newOp.getParam();
                logger.info("{}.param0={}", newOp.getName(), opParam);
                // 删除多余的角色
                for (String r : roles) {
                    if (opParam.containsKey(r)) {
                        opParam.remove(r);
                    }
                }
                logger.info("{}.param1={}", newOp.getName(), opParam);

                // 修改合并指定角色
                if (opParam.containsKey(role) == true) {
                    ObjectMapper mapper = new ObjectMapper();
                    if (opParam.get(role).getClass() == ArrayList.class) {
                        List<Object> list = mapper.convertValue(opParam.get(role), List.class);
                        if (list == null || list.isEmpty() == true) {
                            ;
                        } else {
                            if (list.get(0).getClass() == LinkedHashMap.class) {
                                int index = 0;
                                for (; index < list.size(); index++) {
                                    Map<String, Object> roleParam = mapper.convertValue(list.get(index), HashMap.class);
                                    if (roleParam == null) {
                                        continue;
                                    }
                                    String id = (String) roleParam.getOrDefault("name", "");
                                    if (id.equals(nodeId)) {
                                        break;
                                    }
                                }
                                if (index < list.size()) {
                                    HashMap<String, Object> roleParam = mapper.convertValue(list.get(index), HashMap.class);
                                    HashMap<String, Object> mergeParam = (HashMap<String, Object>) CommonUtils.merge(opParam, roleParam);
                                    // merge之后删除name字段
                                    if (mergeParam.containsKey("name")) {
                                        mergeParam.remove("name");
                                    }
                                    opParam = mergeParam;
                                }
                            }
                        }
                    }
                }
                // 前面已经合并了roleParam; 删除当前角色的roleParam
                opParam.remove(role);
                newOp.setParam(opParam);
                return newOp;
            }
        }
        return null;
    }

    // 输入局部定义的dag覆盖全局dag
    public static DAG merge(DAG baseDag, DAG owDag) throws JsonProcessingException, DAGException {
        // 检查基础dag配置
        DAG dag = DAG.fromJson(baseDag.toJson());
        dag.check();

        // 覆盖dag.resource: 如果confOw不为空则覆盖基础配置
        Map<String, Object> resource = dag.getResource();
        Map<String, Object> owResource = owDag.getResource();
        if (owResource == null) {
            DAGErrorEnums error = DAGErrorEnums.OVERWRITE_RESOURCE_FIELD_EMPTY;
            logger.error("error={}", error.getMessage());
        } else {
            Map<String, Object> mergeResource = CommonUtils.merge(resource, owResource);
            dag.setResource(mergeResource);
        }

        // 覆盖dag.operator; 增量更新
        List<DAGOperator> operator = dag.getOperator();
        List<DAGOperator> owOperator = owDag.getOperator();
        if (owOperator == null || owOperator.isEmpty() == true) {
            DAGErrorEnums error = DAGErrorEnums.OVERWRITE_OPERATOR_LIST_EMPTY;
            logger.error("error={}", error.getMessage());
            throw new DAGException(error);
        }
        for (DAGOperator owOp : owOperator) {
            String name = owOp.getName();
            if (dag.getName2Index().containsKey(name) == false) {
                DAGErrorEnums error = DAGErrorEnums.OVERWRITE_OPERATOR_NAME_NOT_MATCH;
                logger.error("{}.name=not in dag name2, error={}", name, error.getMessage());
                throw new DAGException(error);
            }
            Integer index = dag.getName2Index().get(name);
            DAGOperator op = operator.get(index);
            DAGOperator mergeOp = DAGOperator.merge(op, owOp);
            operator.set(index, mergeOp);
        }
        return dag;
    }

    public static DAG merge(String baseConf, String owConf) throws FileNotFoundException, DAGException, JsonProcessingException {
        // 解析两份dag配置
        Yaml yaml = new Yaml(new Constructor(DAG.class));
        InputStream baseIns = new FileInputStream(new File(baseConf));
        DAG baseDag = yaml.load(baseIns);
        InputStream owIns = new FileInputStream(new File(owConf));
        DAG owDag = yaml.load(owIns);
        DAG dag = DAG.merge(baseDag, owDag);
        return dag;
    }

    // json字符串转换为结构体
    public static DAG fromJson(String json) throws JsonProcessingException {
        ObjectMapper mapper = new ObjectMapper();
        DAG obj = mapper.readValue(json, DAG.class);
        return obj;
    }

    // 结构体转换为json字符串
    public String toJson() throws JsonProcessingException {
        Map<String, Object> map = new HashMap<>();
        map.put("role", role);
        map.put("resource", resource);
        map.put("operator", operator);
        ObjectMapper mapper = new ObjectMapper();
        String json = mapper.writeValueAsString(map);
        return json;
    }

}


